package com.zhuqian.myapplication

import android.graphics.Rect
import org.pytorch.IValue
import java.util.*
import kotlin.math.max
import kotlin.math.min

object Processor {
    private fun to_2d_array(
        array: FloatArray,
        class_num: Int,
        type: String,
    ): Array<FloatArray> {
        return if (type == "yolov5") {
            val cols = class_num + 5
            val rows = array.size / cols
            val res = Array(rows) { FloatArray(cols) }
            for (i in 0 until rows) {
                for (j in 0 until cols) {
                    res[i][j] = array[i * cols + j]
                }
            }
            res
        } else {
            val cols = class_num + 4
            val rows = array.size / cols
            val res = Array(rows) { FloatArray(cols) }
            for (i in 0 until rows) {
                for (j in 0 until cols) {
                    res[i][j] = array[j * rows + i]
                }
            }
            res
        }
    }

    fun parse_output(
        model_output: IValue,
        class_num: Int,
        scaleX: Float, scaleY: Float,
        limit: Int, iou: Float, conf: Float,
        type: String
    ): ArrayList<Result> {

        val outputs = if (type == "yolov5") {
            model_output.toTuple()[0].toTensor().dataAsFloatArray
        } else {
            model_output.toTensor().dataAsFloatArray
        }
        val output2d = to_2d_array(outputs, class_num, type = type)
        println(Arrays.toString(output2d[0]))
        println(Arrays.toString(output2d[1]))
        val results = ArrayList<Result>()
        val class_offset = if (type == "yolov5") 5 else 4
        // 解析预测结果
        for (line in output2d) {
            val x = line[0]
            val y = line[1]
            val w = line[2]
            val h = line[3]
            val left = (x - w / 2) * scaleX
            val top = (y - h / 2) * scaleY
            val right = left + w * scaleX
            val bottom = top + h * scaleY
            val rect = Rect(
                left.toInt(),
                top.toInt(),
                right.toInt(),
                bottom.toInt()
            )
            var classIndex = 0
            if (class_num > 1) {
                var max = line[class_offset]
                for (j in 0 until class_num) {
                    if (line[class_offset + j] > max) {
                        max = line[class_offset + j]
                        classIndex = j
                    }
                }
            }
            val score = if (type == "yolov5") line[4] else line[4 + classIndex]
            val result = Result(classIndex, score, rect)
            results.add(result)
        }
        // 非极大值抑制
        val ret = NMS(results, limit, iou, conf)
        for (result in ret) {
            println(result)
        }
        return ret
    }


    private fun IOU(a: Rect, b: Rect): Float {
        val areaA = ((a.right - a.left) * (a.bottom - a.top)).toFloat()
        if (areaA <= 0.0) return 0.0f
        val areaB = ((b.right - b.left) * (b.bottom - b.top)).toFloat()
        if (areaB <= 0.0) return 0.0f
        val intersectionMinX = max(a.left.toDouble(), b.left.toDouble()).toFloat()
        val intersectionMinY = max(a.top.toDouble(), b.top.toDouble()).toFloat()
        val intersectionMaxX = min(a.right.toDouble(), b.right.toDouble()).toFloat()
        val intersectionMaxY = min(a.bottom.toDouble(), b.bottom.toDouble()).toFloat()
        val intersectionArea = (max(
            (intersectionMaxY - intersectionMinY).toDouble(),
            0.0
        ) * max((intersectionMaxX - intersectionMinX).toDouble(), 0.0)).toFloat()
        return intersectionArea / (areaA + areaB - intersectionArea)
    }

    private fun NMS(boxes: List<Result>, limit: Int, iou: Float, conf: Float): ArrayList<Result> {
        // 置信度过滤
        val filterBox = boxes
            .filter { it.score > conf }          // 置信度过滤
            .sortedByDescending { it.score }     // 按置信度排序
        val selected = ArrayList<Result>()
        val active = BooleanArray(boxes.size)
        Arrays.fill(active, true)
        var numActive = active.size

        // 算法很简单：从具有最高分数的框开始。
        // 删除任何剩余的框，这些框与它的重叠超过给定的阈值
        // 量。如果还有剩余的框（即这些框与任何框都不重叠）
        // 先前的框），然后重复此过程，直到没有更多的框
        // 或已达到数量上限。
        var done = false
        var i = 0
        while (i < filterBox.size && !done) {
            if (active[i]) {
                val boxA: Result = filterBox.get(i)
                selected.add(boxA)
                if (selected.size >= limit) break           // 限制数量
                for (j in i + 1 until filterBox.size) {
                    if (active[j]) {
                        val (_, _, rect) = filterBox.get(j)
                        if (IOU(boxA.rect, rect) > iou) {  // IOU阈值
                            active[j] = false
                            numActive -= 1
                            if (numActive <= 0) {
                                done = true
                                break
                            }
                        }
                    }
                }
            }
            i++
        }
        return selected
    }
}